AWS Lambdaから別アカウントのAmazon Bedrockを呼び出してみた

平井裕二

2024.11.01

はじめに

本記事では、AWS Lambdaから別アカウントのAmazon Bedrockを呼び出す方法について解説します。
既存のシステムとは別のAWSアカウントでAmazon Bedrockを利用する必要がある場合に、この方法が有用です。

アカウント構成とそれぞれのアカウントで作成するリソースは以下の通りです。

アカウントA：呼び出し元（Lambda側アカウント）
- Lambda：Amazon Bedrockを呼び出す
- IAMロール：Lambda実行用のロール（AssumeRole権限を含む）
アカウントB：呼び出し先（Bedrock側アカウント）
- Amazon Bedrock：利用可能にする
- IAMロール：Amazon Bedrockへのアクセス権限を持つロール

処理の流れは以下の通りです。

LambdaがAWS Security Token Service（STS）を使用してBedrock側アカウントのIAMロールを引き受けます
取得した一時認証情報を使用してBedrockにアクセスします
Bedrockからのレスポンスを受け取ります

前提条件

Amazon Bedrockで利用したいモデルのアクセスが有効化されていること

Amazon Bedrockのモデルアクセス設定画面

Lambda側アカウントでの作業

Lambda側アカウントでは、以下の2つのリソースを作成します

Lambda実行用のIAMロール（AssumeRole権限含む）
Bedrock呼び出し用のLambda

Lambda用のIAMポリシーとIAMロール作成

まず、Lambda用のIAMロールを作成します。このロールには、Lambdaの基本実行権限とBedrock側のロールを引き受ける権限が必要です。

AWS CloudShellで以下のコマンドを実行し、Lambda用のIAMロールを作成します。

$ aws iam create-role \
    --role-name bedrock-cross-account-invoke-role \
    --assume-role-policy-document '{
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {
                    "Service": "lambda.amazonaws.com"
                },
                "Action": "sts:AssumeRole"
            }
        ]
    }'

Lambda用のIAMロールに、Bedrock側アカウントのロールを引き受けるためのAssumeRole権限を持つポリシーを作成し、アタッチします。

${BEDROCK_ACCOUNT_ID}は、Bedrockを利用するAWSアカウントのIDに置き換えてください。

$ aws iam create-policy \
    --policy-name lambda-bedrock-assumerole-policy \
    --policy-document '{
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": "sts:AssumeRole",
                "Resource": "arn:aws:iam::${BEDROCK_ACCOUNT_ID}:role/bedrock-cross-account-role"
            }
        ]
    }'

AssumeRoleポリシーをアタッチします。

${LAMBDA_ACCOUNT_ID}は、Lambdaを利用するAWSアカウントのIDに置き換えてください。

$ aws iam attach-role-policy \
    --role-name bedrock-cross-account-invoke-role \
    --policy-arn arn:aws:iam::${LAMBDA_ACCOUNT_ID}:policy/lambda-bedrock-assumerole-policy

Lambdaの基本実行ポリシーであるAWSLambdaBasicExecutionRoleをIAMロールにアタッチします。

$ aws iam attach-role-policy \
    --role-name bedrock-cross-account-invoke-role \
    --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

Lambdaを作成

Lambdaは、マネジメントコンソール上で作成します。

bedrock-cross-account-invokeという名前で作成します。
IAMロールは先程作成したbedrock-cross-account-invoke-roleを選択します。
${BEDROCK_ACCOUNT_ID}は、Bedrockを利用するAWSアカウントのIDに置き換えてください。

import boto3
import json

BEDROCK_ROLE_ARN = 'arn:aws:iam::${BEDROCK_ACCOUNT_ID}:role/bedrock-cross-account-role'
MODEL_ID = 'anthropic.claude-3-5-sonnet-20240620-v1:0'
REGION_NAME = 'ap-northeast-1'

def get_bedrock_client():
    sts = boto3.client('sts')
    role = sts.assume_role(
        RoleArn=BEDROCK_ROLE_ARN,
        RoleSessionName="assume-role"
    )

    return boto3.client(
        'bedrock-runtime',
        aws_access_key_id=role['Credentials']['AccessKeyId'],
        aws_secret_access_key=role['Credentials']['SecretAccessKey'],
        aws_session_token=role['Credentials']['SessionToken'],
        region_name=REGION_NAME
    )

def invoke_bedrock(client, prompt):
    response = client.converse(
        modelId=MODEL_ID,
        messages=[{
            "role": "user",
            "content": [{
                "text": prompt
            }]
        }],
        inferenceConfig={
            "maxTokens": 1000,
            "temperature": 0
        }
    )

    response_text = response['output']['message']['content'][0]['text']
    print('Received response_text:' + json.dumps(response_text, ensure_ascii=False))
    return response_text

def lambda_handler(event, context):
    client = get_bedrock_client()
    prompt = '1+1は？答えだけ教えて'
    bedrock_response = invoke_bedrock(client, prompt)

    return {
        'statusCode': 200,
        'body': bedrock_response
    }

Lambdaでは、以下の3つの主要な処理を行います

STSを使用してBedrock側のIAMロールを引き受ける
- STSを使用して一時的な認証情報を取得
- 指定したIAMロールARNに対してAssumeRoleを実行
取得した認証情報でBedrockクライアントを作成
- 一時的な認証情報を使用してBedrockクライアント初期化
- 指定したリージョンでクライアントを設定
Bedrockのモデルを呼び出す
- Claude 3.5 Sonnetモデルを使用してテキスト生成を実行
- 設定したパラメータ（maxTokens、temperature）に基づいて推論を実行

Bedrock側アカウントでの作業

Bedrock側アカウントでは、Lambda側アカウントのIAMロールが引き受けることができるIAMロールを作成します。

以下のコマンドを実行します。

$ aws iam create-role \
    --role-name bedrock-cross-account-role \
    --assume-role-policy-document '{
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {
                    "AWS": "arn:aws:iam::${LAMBDA_ACCOUNT_ID}:role/bedrock-cross-account-invoke-role"
                },
                "Action": "sts:AssumeRole"
            }
        ]
    }'

このIAMロールには、Bedrockへのアクセス権限が必要です。
AmazonBedrockFullAccessをアタッチします。

$ aws iam attach-role-policy \
    --role-name bedrock-cross-account-role \
    --policy-arn arn:aws:iam::aws:policy/AmazonBedrockFullAccess

動作確認

Lambdaの動作確認を行います。

Lambdaをテスト実行
- マネジメントコンソールからテスト実行を行います
- テストイベントの内容は空のJSONで構いません
実行結果
以下のようなレスポンスが返ってきました

{
  "statusCode": 200,
  "body": "2"
}

もしLambda側アカウントのリソース（例えばS3バケット）にリクエストしたい場合、権限を追加した上で、以下のコードを追加します。

この部分では、デフォルトの認証情報（Lambda実行ロール）を使用してS3クライアントを作成しているため、Lambda側のアカウントのS3バケットにアクセスします。

def list_s3_objects():
    s3_client = boto3.client('s3')  # デフォルトの認証情報（Lambda実行ロール）を使用
    response = s3_client.list_objects_v2(
        Bucket='cm-hirai-bedrock'
    )
    return [obj['Key'] for obj in response.get('Contents', [])]

~中略~
    s3_objects = list_s3_objects()

全コード（クリックで展開）

import boto3
import json

BEDROCK_ROLE_ARN = 'arn:aws:iam::${BEDROCK_ACCOUNT_ID}:role/bedrock-cross-account-role'
MODEL_ID = 'anthropic.claude-3-5-sonnet-20240620-v1:0'
REGION_NAME = 'ap-northeast-1'

def get_bedrock_client():
    sts = boto3.client('sts')
    role = sts.assume_role(
        RoleArn=BEDROCK_ROLE_ARN,
        RoleSessionName="assume-role"
    )

    return boto3.client(
        'bedrock-runtime',
        aws_access_key_id=role['Credentials']['AccessKeyId'],
        aws_secret_access_key=role['Credentials']['SecretAccessKey'],
        aws_session_token=role['Credentials']['SessionToken'],
        region_name=REGION_NAME
    )

def invoke_bedrock(client, prompt):
    response = client.converse(
        modelId=MODEL_ID,
        messages=[{
            "role": "user",
            "content": [{
                "text": prompt
            }]
        }],
        inferenceConfig={
            "maxTokens": 1000,
            "temperature": 0
        }
    )

    response_text = response['output']['message']['content'][0]['text']
    print('Received response_text:' + json.dumps(response_text, ensure_ascii=False))
    return response_text

def list_s3_objects():
    s3_client = boto3.client('s3')
    response = s3_client.list_objects_v2(
        Bucket='cm-hirai-bedrock'
    )
    return [obj['Key'] for obj in response.get('Contents', [])]

def lambda_handler(event, context):
    client = get_bedrock_client()
    prompt = '1+1は？答えだけ教えて'
    bedrock_response = invoke_bedrock(client, prompt)

    # 処理例：S3の処理
    s3_objects = list_s3_objects()

    return {
        'statusCode': 200,
        'body': {
            'bedrock_response': bedrock_response,
            's3_objects': s3_objects
        }
    }

対して、Bedrockの場合は、明示的にクロスアカウントの認証情報を取得しています。

def get_bedrock_client():
    sts = boto3.client('sts')
    role = sts.assume_role(
        RoleArn=BEDROCK_ROLE_ARN,
        RoleSessionName="assume-role"
    )